import pandas as pd
import numpy as np
from scipy.fft import fft
from scipy.signal import stft
import pywt
import matplotlib.pyplot as plt

from DWDataReader import DWDataReaderExample


# 基本统计量
def compute_basic(df):
    time_stats = df.groupby('channel_name').agg(
        mean=('Value', 'mean'),
        max=('Value', 'max'),
        min=('Value', 'min'),
        peak_to_peak=('Value', lambda x: x.max() - x.min())
    )
    return time_stats

def compute_time_features(df):
    """计算时域统计特征"""
    # 基本统计量
    time_stats = compute_basic(df)

    # 过零率计算（添加阈值避免噪声误判）
    threshold = 1e-8  # 根据实际信号调整阈值
    df['sign'] = np.sign(df['Value'])
    df['crossing'] = (df['sign'] * df['sign'].shift(1) < 0) & (np.abs(df['Value']) > threshold)
    zcr = df.groupby('channel_name')['crossing'].mean()

    time_stats['zero_crossing_rate'] = zcr
    return time_stats


def compute_fft_features(group, sampling_rate=10000):
    """计算频域特征"""
    values = group['Value'].values
    n = len(values)
    if n < 2:
        return None  # 跳过数据不足的通道

    fft_values = fft(values)
    freqs = np.fft.fftfreq(n, d=1 / sampling_rate)
    psd = (np.abs(fft_values) ** 2) / n
    dominant_freq = freqs[np.argmax(psd)]

    return {
        'dominant_frequency': dominant_freq,
        'psd': psd.tolist()
    }


def compute_time_freq_features(group, sampling_rate=10000):
    """计算时频域特征"""
    values = group['Value'].values
    n = len(values)
    if n < 256:  # STFT 需要足够的数据长度
        return None

    # STFT
    f, t, Z = stft(values, fs=sampling_rate, nperseg=256, noverlap=128)
    stft_psd = np.abs(Z) ** 2

    # 小波变换
    max_level = min(int(np.floor(np.log2(n))), 5) if n > 1 else 0  # 限制最大层数
    coeffs = pywt.wavedec(values, 'db4', level=max_level)
    energy = [np.sum(np.abs(c) ** 2) for c in coeffs]
    scales = [2 ** (level + 1) for level in range(max_level)]  # 生成尺度列表

    return {
        'stft_psd': stft_psd.tolist(),
        'wavelet_energy': energy,
        'wavelet_scales': scales
    }


def main():
    # 加载数据
    data = DWDataReaderExample.get_dxd_data_sample()

    # 转换为 DataFrame
    transformed_data = []
    for channel, samples in data.items():
        for sample in samples:
            transformed_data.append({
                "Time": sample["Time"],
                "Value": sample["Value"],
                "channel_name": channel
            })
    df = pd.DataFrame(transformed_data)

    # 时域特征
    time_features = compute_time_features(df)
    print("时域特征:\n", time_features)

    # 频域特征
    fft_results = {}
    for channel, group in df.groupby('channel_name'):
        result = compute_fft_features(group)
        if result:
            fft_results[channel] = result
    print("\n频域特征:\n", fft_results)

    # 时频域特征
    time_freq_results = {}
    for channel, group in df.groupby('channel_name'):
        result = compute_time_freq_features(group)
        if result:
            time_freq_results[channel] = result

    # ======================== 可视化部分 ========================
    # 1. 时域信号波形图
    plt.figure(figsize=(12, 6))
    for channel in df['channel_name'].unique()[:3]:  # 仅展示前3个通道避免拥挤
        channel_df = df[df['channel_name'] == channel]
        plt.plot(channel_df['Time'], channel_df['Value'], label=channel)
    plt.title("Time Domain Waveform")  # 英文标题
    plt.xlabel("Time (s)")  # 英文坐标标签
    plt.ylabel("Amplitude")
    plt.legend()
    plt.grid(True)
    plt.show()

    # 2. 时域统计量柱状图
    plt.figure(figsize=(12, 6))
    time_features[['mean', 'max', 'min', 'peak_to_peak']].plot(kind='bar')
    plt.title("Time Domain Statistics")  # 英文标题
    plt.xlabel("Channel")
    plt.ylabel("Value")
    plt.xticks(rotation=45)
    plt.grid(True)
    plt.show()

    # 3. 频域主频对比图
    if fft_results:
        dominant_freqs = {k: v['dominant_frequency'] for k, v in fft_results.items()}
        plt.figure(figsize=(12, 6))
        plt.bar(dominant_freqs.keys(), dominant_freqs.values())
        plt.title("Dominant Frequency Comparison")
        plt.xlabel("Channel")
        plt.ylabel("Frequency (Hz)")
        plt.xticks(rotation=45)
        plt.grid(True)
        plt.show()

    # 4. 频域功率谱密度图
    if fft_results:
        first_channel = list(fft_results.keys())[0]
        freqs = np.array(fft_results[first_channel]['psd'])
        plt.figure(figsize=(12, 6))
        plt.plot(freqs[:len(freqs) // 2])  # 显示正频率部分
        plt.title(f"Power Spectrum - {first_channel}")
        plt.xlabel("Frequency (Hz)")
        plt.ylabel("Power")
        plt.grid(True)
        plt.show()

    # ======================== 时频域可视化 ========================
    # 5. STFT 时频图
    if time_freq_results:
        first_channel = list(time_freq_results.keys())[0]
        group = df[df['channel_name'] == first_channel]
        values = group['Value'].values

        f, t, Z = stft(values, fs=10000, nperseg=256, noverlap=128)
        plt.figure(figsize=(12, 4))
        plt.pcolormesh(t, f, np.abs(Z) ** 2, shading='gouraud', cmap='viridis')
        plt.title(f'STFT Spectrum - {first_channel}')
        plt.ylabel('Frequency (Hz)')
        plt.xlabel('Time (s)')
        plt.colorbar(label='Power')
        plt.show()

    # 6. 小波能量分布
    if time_freq_results:
        first_channel = list(time_freq_results.keys())[0]
        energy = time_freq_results[first_channel]['wavelet_energy']
        plt.figure(figsize=(12, 4))
        plt.plot(range(len(energy)), energy, 'o-')
        plt.title(f'Wavelet Energy - {first_channel}')
        plt.xlabel('Scale Level')
        plt.ylabel('Energy')
        plt.grid(True)
        plt.show()


if __name__ == "__main__":
    main()